#!/bin/env python
# -*- coding: utf-8 -*-

"""
File Name: zk_util.py
Author: morre <morre@vip.qq.com>
Create Time: 2022/12/31 17:08:04
Brief:
"""

from __future__ import absolute_import
import struct
import time
import logging
from google.protobuf import message
from kazoo.client import KazooClient
from kazoo.handlers.threading import KazooTimeoutError
from kazoo.exceptions import NoNodeError

from hbase_client.pb.ZooKeeper_pb2 import MetaRegionServer
from .hbase_exceptions import ZookeeperResponseException, ZookeeperConnectionException, ZookeeperZNodeException


logger = logging.getLogger("hbase_client." + __name__)
# logger.setLevel(logging.DEBUG)

# LocateMetaRegion takes a string representing the location of the ZooKeeper
# quorum. It then asks ZK for the location of the MetaRegionServer,
# returning a tuple containing (host_name, port).

# default: establish_connection_timeout=5, missing_znode_retries=5
#  TODO: return full info.


def locate_meta_region(
    zkquorum, establish_connection_timeout, missing_znode_retries,
    zk_root_path):
    # MetaRegionServer information is located at /hbase/meta-region-server
    if not zk_root_path:
        zk_root_path = "/hbase"
    zkpath = zk_root_path + "/meta-region-server"
    rsp = zk_request(
        zkquorum, zkpath, establish_connection_timeout, missing_znode_retries)
    # parse MetaRegionServer message
    meta = MetaRegionServer()
    try:
        meta.ParseFromString(rsp)
    except message.DecodeError as exc:
        raise ZookeeperResponseException(
            "Failed to parse MetaRegionServer from response."
        ) from exc
    logger.info("Discovered MetaRegionServer at %s:%s",
                meta.server.host_name, meta.server.port)
    return meta.server.host_name, int(meta.server.port)


def zk_request(
    zkquorum, zkpath, establish_connection_timeout, missing_znode_retries):
    # Using Kazoo for interfacing with ZK
    zk = KazooClient(hosts=zkquorum)
    try:
        zk.start(timeout=establish_connection_timeout)
    except KazooTimeoutError as exc:
        raise ZookeeperConnectionException(
            f"Cannot connect to ZooKeeper at {zkquorum}") from exc

    rsp, znodestat = None, None
    for _ in range(missing_znode_retries + 1):
        try:
            rsp, znodestat = zk.get(zkpath)
        except NoNodeError:
            logger.warning(
                "ZooKeeper does not contain meta-region-server node. "
                "Retrying in 2 seconds. (%s retries remaining)",
                missing_znode_retries)
            time.sleep(2.0)
            continue
        else:
            break
    _ = znodestat
    if rsp is None:
        raise ZookeeperZNodeException(
            f"ZooKeeper does not contain node in path: {zkpath}.")
    # We don"t need to maintain a connection to ZK. If we need it again we"ll
    # recreate the connection. A possible future implementation can subscribe
    # to ZK and listen for when RegionServers go down, then pre-emptively
    # reestablish those regions instead of waiting for a failed rpc to come
    # back. Only issue is that if too many clients subscribe ZK may become
    # overloaded.
    zk.stop()
    # the message contains at least 5 bytes with the following structure:
    # (1B)(4B)... => (b"\xff")(payload_size)...
    if len(rsp) < 5:
        raise ZookeeperResponseException(
            f"ZooKeeper returned too few response. Response size: {len(rsp)}.")
    # The first byte must be \xff and the next four bytes are a little-endian
    # uint32 containing the length of the payload.
    first_byte, payload_size = struct.unpack(">cI", rsp[:5])
    if first_byte != b"\xff":
        # Malformed response
        raise ZookeeperResponseException(
            "ZooKeeper returned an invalid response, b\"\\xff\" excepted, "
            f"got {first_byte}")
    if payload_size < 1 or payload_size > 65000:
        # Is this really an error?
        raise ZookeeperResponseException(
            f"ZooKeeper returned invaild payload size {payload_size}.")
    # ZNode data in HBase are serialized protobufs with a four byte magic
    # "PBUF" prefix.
    magic = struct.unpack(">4s", rsp[payload_size + 5:payload_size + 9])[0]
    if magic != b"PBUF":
        # 4 bytes: PBUF
        raise ZookeeperResponseException(
            "ZooKeeper returned an invalid response. b\"PBUF\" expected, "
            f"got {magic}.")
    return rsp[payload_size + 9:]
